import numpy as np
from PIL import Image
import caffe
import logging
import os
import scipy

FORMAT = '[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
DATEFORMAT = '%H:%M:%S'
logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
logging.debug('start')

home_dir = os.getenv('HOME', '')
db_root = home_dir + '/data'
out_root = db_root + '/out'
flist_name = out_root + '/filelist'
lesion_root = home_dir + '/data/pixellevel/出血-红色'
# original_root = home_dir + '/data/pixellevel/原始'
original_root = './'

out_original_root = out_root + '/original'
out_lesion_root = out_root + '/ha'
home_dir = os.getenv('HOME', '')
#
#
#
# # load image, switch to BGR, subtract mean, and make dims C x H x W for Caffe
# im = Image.open('pascal/VOC2010/JPEGImages/2007_000129.jpg')
# in_ = np.array(im, dtype=np.float32)
# in_ = in_[:,:,::-1]
# in_ -= np.array((104.00698793,116.66876762,122.67891434))
# in_ = in_.transpose((2,0,1))
#
# # load net
# net = caffe.Net('voc-fcn8s/deploy.prototxt', 'voc-fcn8s/fcn8s-heavy-pascal.caffemodel', caffe.TEST)
# # shape for input (data blob is N x C x H x W), set data
# net.blobs['data'].reshape(1, *in_.shape)
# net.blobs['data'].data[...] = in_
# # run net and take argmax for prediction
# net.forward()
# out = net.blobs['score'].data[0].argmax(axis=0)


caffe.set_mode_gpu()
caffe.set_device(0)

net = caffe.Net('deploy.prototxt',
                'snapshot/train_iter_1600.caffemodel',
                caffe.TEST)

out_dir = 'out/'
if not os.path.exists(out_dir):
    os.mkdir(out_dir)

# flist = os.listdir(original_root)
flist = ['test.jpg']
index = 0
for n in flist:
    index += 1
    if n.find('-') >= 0:
        continue
    imgname = original_root + '/' + n
    im = Image.open(imgname)
    in_ = np.array(im, dtype=np.float32)
    in_ = in_[:, :, ::-1]
    in_ -= np.array((149.54685784, 88.26339913, 37.44157289))
    in_ = in_.transpose((2, 0, 1))
    # in_ /= 255

    # load net
    # net = caffe.Net('voc-fcn8s/deploy.prototxt', 'voc-fcn8s/fcn8s-heavy-pascal.caffemodel', caffe.TEST)
    # shape for input (data blob is N x C x H x W), set data
    net.blobs['data'].reshape(1, *in_.shape)
    net.blobs['data'].data[...] = in_
    # run net and take argmax for prediction
    net.forward()
    out = net.blobs['score'].data[0]
    print(out.shape)

    # color = True
    # img = data.astype(np.float32)
    # if img.ndim == 2:
    #     img = img[:, :, np.newaxis]
    #     if color:
    #         img = np.tile(img, (1, 1, 3))
    # elif img.shape[2] == 4:
    #     img = img[:, :, :3]
    #
    # print('handling', imgname, '-->', index, sep=' ', end='\n')
    # net.blobs['data'].data[...] = transformer.preprocess('data', img)
    # net.forward()
    # out = net.blobs['score'].data[0].argmax(axis=0)
    # print(type(out))
    # print(out.shape)
    scipy.misc.imsave(out_dir + '%010d_1.jpg' % index, out[1, :, :])
    scipy.misc.imsave(out_dir + '%010d_0.jpg' % index, out[0, :, :])
    scipy.misc.imsave(out_dir + '%010d_2.jpg' % index, out.argmax(axis=0))
    # print(n, out, sep='\t')
